package pkg

import (
	"flag"
	"fmt"
	"log"
	"net/url"
	"os"
	"runtime"
	"slices"
	"strconv"
	"strings"
	"text/tabwriter"

	"github.com/fatih/color"
	"golang.org/x/time/rate"
)

var (
	version   string
	useragent string

	generalOptions  []FlagStruct
	generateOptions []FlagStruct
	requestOptions  []FlagStruct
	crawlOptions    []FlagStruct
	wordlistOptions []FlagStruct
)

type FlagStruct struct {
	LongFlag    string
	ShortFlag   string
	Description string
}

func init() {

}

func ParseFlags(vers string) {
	/* Getting Command-line flags */
	version = vers
	useragent = "WebCacheVulnerabilityScanner v" + version
	pathPrefix := ""
	if runtime.GOOS == "windows" {
		pathPrefix = "C:"
	}

	// General Options
	techniqueNames := "deception,cookies,css,forwarding,smuggling,dos,headers,parameters,fatget,cloaking,splitting,pollution,encoding"

	var ignoreStatus string

	appendInt(&generalOptions, &Config.Verbosity,
		"verbosity", "v", 1, "Set verbosity. 0 = quiet, 1 = normal, 2 = verbose")
	appendFloat(&generalOptions, &Config.ReqRate,
		"reqrate", "rr", float64(rate.Inf), "Requests per second. Float value. Has to be greater than 0. Default value is infinite")
	appendInt(&generalOptions, &Config.Threads,
		"threads", "t", 20, "Threads to use. Default value is 20")
	appendInt(&generalOptions, &Config.TimeOut,
		"timeout", "to", 15, "Seconds until timeout. Default value is 15")
	appendString(&generalOptions, &Config.OnlyTest,
		"onlytest", "ot", "", "Choose which tests to run. Use the , separator to specify multiple ones. Example: -onlytest '"+techniqueNames+"'")
	appendString(&generalOptions, &Config.SkipTest,
		"skiptest", "st", "", "Choose which tests to not run. Use the , separator to specify multiple ones. Example: -skiptest '"+techniqueNames+"'")
	appendBoolean(&generalOptions, &Config.UseProxy,
		"useproxy", "up", false, "Do you want to use a proxy?")
	appendString(&generalOptions, &Config.ProxyURL,
		"proxyurl", "purl", "http://127.0.0.1:8080", "Url for the proxy. Default value is http://127.0.0.1:8080")
	appendBoolean(&generalOptions, &Config.Force,
		"force", "f", false, "Perform the tests no matter if there is a cache or even the cachebuster works or not")
	appendString(&generalOptions, &ignoreStatus,
		"ignorestatus", "is", "", "Ignore a specific status code for cache poisoning")
	appendString(&generalOptions, &Config.ReasonTypes,
		"reasontypes", "rt", "body,header,status,length", "Choose which reason types to use for cache poisoning. Choose from: body (reflection in body),header (reflection in header), status (change of status code), length (change of body length). Default is 'body,header,status,length'")
	appendInt(&generalOptions, &Config.CLDiff,
		"contentlengthdifference", "cldiff", 5000, "Threshold for reporting possible Finding, when 'poisoned' response differs more from the original length. Default is 5000. 0 = don't check. May be prone to false positives!")
	appendInt(&generalOptions, &Config.HMDiff,
		"hitmissdifference", "hmdiff", 30, "Threshold for time difference (milliseconds) between cache hit and cache miss responses. Default is 30")
	appendBoolean(&generalOptions, &Config.SkipTimebased,
		"skiptimebased", "stime", false, "Skip checking if a repsonse gets cached by measuring time differences (may be prone to false positives, or increase hitmissdifference)")
	appendBoolean(&generalOptions, &Config.SkipWordlistCachebuster,
		"skipwordlistcachbuster", "swordlistcb", false, "Skip using wordlists to check for cachebusters (may be time intensive)")
	appendString(&generalOptions, &Config.CacheHeader,
		"cacheheader", "ch", "", "Specify a custom cache header")
	appendBoolean(&generalOptions, &Config.DisableColor,
		"nocolor", "nc", false, "Disable colored output")
	appendBoolean(&generalOptions, &Config.DisableStatusLine,
		"nostatusline", "ns", false, "Disable status line output")

	// Generate Options
	appendString(&generateOptions, &Config.GeneratePath,
		"generatepath", "gp", "./", "Path all files (log, report, completed) will be written to. Example: -gp '"+pathPrefix+"/p/a/t/h/'. Default is './'")
	appendBoolean(&generateOptions, &Config.GenerateReport,
		"generatereport", "gr", false, "Do you want a report to be generated?")
	appendBoolean(&generateOptions, &Config.EscapeJSON,
		"escapejson", "ej", false, "Do you want HTML special chars to be encoded in the report?")
	appendBoolean(&generateOptions, &Config.GenerateCompleted,
		"generatecompleted", "gc", false, "Do you want a list with completed URLs to be generated?")
	appendBoolean(&generateOptions, &Config.GenerateLog,
		"generatelog", "gl", false, "Do you want a log file to be created?")

	// Request Options
	var (
		urlStr           string
		setCookiesStr    string
		setHeadersStr    string
		setParametersStr string
		setBodyStr       string
		userAgentChrome  bool
	)

	appendString(&requestOptions, &urlStr,
		"url", "u", "", "Url to scan. Has to start with http:// or https://. Otherwise use file: to specify a file with (multiple) urls. E.g. -u https://www.example.com or -u file:templates/url_list")
	appendBoolean(&requestOptions, &Config.UseHTTP,
		"usehttp", "http", false, "Use http instead of https for URLs, which doesn't specify either one")
	appendBoolean(&requestOptions, &Config.DeclineCookies,
		"declineCookies", "dc", false, "Do you don't want to use cookies, which are received in the response of the first request?")
	appendString(&requestOptions, &Config.CacheBuster,
		"cachebuster", "cb", "cbwcvs", "Specify the cachebuster to use. The default value is cbwcvs")
	appendString(&requestOptions, &setCookiesStr,
		"setcookies", "sc", "", "Set a Cookie. Otherwise use file: to specify a file with urls. E.g. -sc uid=123 or -sc file:templates/cookie_list")
	appendString(&requestOptions, &setHeadersStr,
		"setheaders", "sh", "", "Set a Header. Otherwise use file: to specify a file with urls. E.g. -sh 'User-Agent: Safari/1.1' or -sh file:templates/header_list")
	appendString(&requestOptions, &setParametersStr,
		"setparameters", "sp", "", "Set a Query Parameter. Otherwise use file: to specify a file with urls. E.g. -sp user=admin or -sp file:templates/parameter_list")
	appendString(&requestOptions, &setBodyStr,
		"setbody", "sb", "", "Set the requests' body. Otherwise use file: to specify a file with urls. E.g. -sb 'admin=true' or -sh file:templates/body_file")
	appendBoolean(&requestOptions, &Config.DoPost,
		"post", "post", false, "Do a POST request instead of a GET request")
	appendString(&requestOptions, &Config.ContentType,
		"contenttype", "ct", "application/x-www-form-urlencoded", "Set the contenttype for a POST Request. Default is application/x-www-form-urlencoded. If you don't want a content-type to be used at all use -ct ''")
	appendString(&requestOptions, &Config.QuerySeparator,
		"parameterseparator", "ps", "&", "Specify the separator for parameters. The default value is &")
	appendBoolean(&requestOptions, &userAgentChrome,
		"useragentchrome", "uac", false, "Set chrome as User-Agent. Default is "+useragent)

	// Crawl Options
	var (
		recExcludeStr string
		recDomainsStr string
	)

	appendInt(&crawlOptions, &Config.Recursivity,
		"recursivity", "r", 0, "Put (via href or src specified) urls at the end of the queue if the domain is the same. Specify how deep the recursivity shall go. Default value is 0 (no recursivity)")
	appendInt(&crawlOptions, &Config.RecLimit,
		"reclimit", "rl", 0, "Define a limit, how many files shall be checked recursively. Default is 0 (unlimited)")
	appendString(&crawlOptions, &Config.RecInclude,
		"recinclude", "rin", "", "Choose which links should be included. Separate with a space. E.g: -rin '.js .css'")
	appendString(&crawlOptions, &recExcludeStr,
		"recexclude", "rex", "", "Use -cp (-completedpath) or -gc (-generatecompleted) to generate a list of already completed URLs. Use -rex path/to/file so the already completed URLs won't be tested again recursively.")
	appendString(&crawlOptions, &recDomainsStr,
		"recdomains", "red", "", "Define an additional domain which is allowed to be added recursively. Otherwise use file: to specify a file with urls. E.g. -sh 'api.example.com' or -sh file:templates/recdomains_list")

	// Wordlist Options
	appendString(&wordlistOptions, &Config.HeaderWordlist,
		"headerwordlist", "hw", "", "Wordlist for headers to test.")
	appendString(&wordlistOptions, &Config.ParameterWordlist,
		"parameterwordlist", "pw", "", "Wordlist for query parameters to test.")

	flag.CommandLine.Usage = help

	// flags need to be parsed, before they are used
	flag.Parse()

	// Check if color should be disabled
	if Config.DisableColor {
		color.NoColor = true
	}
	fmt.Printf(getLogo()+"\nWCVS - the Web Cache Vulnerability Scanner. (v%s)"+"\n\n", version)

	/* Checking values of Flags */
	if len(flag.Args()) > 0 {
		msg := fmt.Sprintf("%s: Args are not supported! Use flags. Use -h or --help to get a list of all supported flags\n", flag.Args())
		PrintFatal(msg)
	}
	if urlStr == "" {
		msg := "No url specified. Use -url or -u. Use -h or --help to get a list of all supported flags\n"
		PrintFatal(msg)
	}

	// Change User Agent
	if userAgentChrome {
		useragent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/135.0.0.0 Safari/537.36"
	}

	// IgnoreStatus string to int slice
	if ignoreStatus != "" {
		statusSlice := strings.Split(ignoreStatus, ",")
		for _, status := range statusSlice {
			statusInt, err := strconv.Atoi(strings.TrimSpace(status))
			if err != nil {
				fmt.Printf(getLogo()+"\nWCVS - the Web Cache Vulnerability Scanner. (v%s)"+"\n\n", version)
				PrintFatal("Error converting to int: " + err.Error())
			}
			Config.IgnoreStatus = append(Config.IgnoreStatus, statusInt)
		}
	}

	// Read RecExcludeURL(s)
	if recExcludeStr != "" {
		Config.RecExclude = ReadLocalFile(recExcludeStr, "RecExclude")
	}

	// Read URL(s)
	Config.Urls = readFile(urlStr, Config.Urls, "URL")

	// Read RecDomain(s)
	Config.RecDomains = readFile(recDomainsStr, Config.RecDomains, "RecDomain")
	for _, u := range Config.Urls { // add all domains from urls to recdomains
		urlParsed, err := url.Parse(strings.TrimSpace(u))
		if err != nil {
			log.Fatal("Error parsing recdomains: ", err)
		}
		domain := urlParsed.Hostname()
		if !slices.Contains(Config.RecDomains, domain) {
			Config.RecDomains = append(Config.RecDomains, domain)
		}
	}

	// Read Cookie(s)
	Config.Cookies = readFile(setCookiesStr, Config.Cookies, "Cookie")

	// Read Header(s)
	Config.Headers = readFile(setHeadersStr, Config.Headers, "Headers")

	// Read Parameter(s)
	Config.Parameters = readFile(setParametersStr, Config.Parameters, "Parameter")

	/* Read Body */
	if strings.HasPrefix(setBodyStr, "path:") {
		bodySlice := ReadLocalFile(setBodyStr, "Body")
		for _, l := range bodySlice {
			l = strings.TrimSuffix(l, "\r")
			l = strings.TrimSpace(l)
			if strings.HasPrefix(l, "//") || l == "" {
				continue
			}
			Config.Body += l
		}
	} else {
		Config.Body = setBodyStr
	}

	// Set Limiter
	Config.Limiter = rate.NewLimiter(rate.Limit(Config.ReqRate), 1)

	Config.OnlyTest = strings.ToLower(Config.OnlyTest)
	Config.SkipTest = strings.ToLower(Config.SkipTest)
}

func readFile(str string, field []string, name string) []string {
	if strings.HasPrefix(str, "file:") {
		return ReadLocalFile(str, name)
	} else {
		return append(field, str)
	}
}

func help() {
	w := new(tabwriter.Writer)
	w.Init(os.Stdout, 8, 8, 0, '\t', 0)

	fmt.Printf(getLogo()+"\nWCVS - the Web Cache Vulnerability Scanner. (v%s)"+"\n\n", version)

	fmt.Println("Published by Hackmanit under http://www.apache.org/licenses/LICENSE-2.0")
	fmt.Println("Author: Maximilian Hildebrand")
	fmt.Println("Repository: https://github.com/Hackmanit/Web-Cache-Vulnerability-Scanner")
	fmt.Println("Blog Post: https://hackmanit.de/en/blog-en/145-web-cache-vulnerability-scanner-wcvs-free-customizable-easy-to-use")
	fmt.Print("Usage: Web-Cache-Vulnerability-Scanner(.exe) [options]\n\n")

	fmt.Println("General Options:")
	fmt.Fprintf(w, "%s\t%s\t%s\n", "--help", "-h", "Show this help and quit")
	writeToWriter(w, generalOptions)

	fmt.Println("\nGenerate Options:")
	writeToWriter(w, generateOptions)

	fmt.Println("\nRequest Options:")
	writeToWriter(w, requestOptions)

	fmt.Println("\nCrawl Options:")
	writeToWriter(w, crawlOptions)

	fmt.Println("\nWordlist Options:")
	writeToWriter(w, wordlistOptions)

	os.Exit(0)
}

func getLogo() string {
	// source: https://patorjk.com/software/taag/#p=display&f=Slant%20Relief&t=wcvs
	logo := `
__/\\____/\\___/\\_____/\\\\\\\\__/\\\____/\\\__/\\\\\\\\\\_     
 _\/\\\__/\\\\_/\\\___/\\\//////__\//\\\__/\\\__\/\\\//////__    
  _\//\\\/\\\\\/\\\___/\\\__________\//\\\/\\\___\/\\\\\\\\\\_   
   __\//\\\\\/\\\\\___\//\\\__________\//\\\\\____\////////\\\_  
    ___\//\\\\//\\\_____\///\\\\\\\\____\//\\\______/\\\\\\\\\\_ 
     ____\///__\///________\////////______\///______\//////////__`

	logo = strings.ReplaceAll(logo, "_", color.HiRedString("_"))
	return logo
}

func writeToWriter(w *tabwriter.Writer, flagStruct []FlagStruct) {
	for _, ts := range flagStruct {
		fmt.Fprintf(w, "--%s\t-%s\t%s\n", ts.LongFlag, ts.ShortFlag, ts.Description)
	}
	w.Flush()
}

func appendString(options *[]FlagStruct, varString *string, longFlag string, shortFlag string, defaultValue string, description string) {
	flag.StringVar(varString, longFlag, defaultValue, "")
	if shortFlag != longFlag {
		flag.StringVar(varString, shortFlag, defaultValue, "")
	}
	*options = append(*options, FlagStruct{
		LongFlag:    longFlag,
		ShortFlag:   shortFlag,
		Description: description})
}

func appendInt(options *[]FlagStruct, varInt *int, longFlag string, shortFlag string, defaultValue int, description string) {
	flag.IntVar(varInt, longFlag, defaultValue, "")
	if shortFlag != longFlag {
		flag.IntVar(varInt, shortFlag, defaultValue, "")
	}
	*options = append(*options, FlagStruct{
		LongFlag:    longFlag,
		ShortFlag:   shortFlag,
		Description: description})
}

func appendFloat(options *[]FlagStruct, varFloat *float64, longFlag string, shortFlag string, defaultValue float64, description string) {
	flag.Float64Var(varFloat, longFlag, defaultValue, "")
	if shortFlag != longFlag {
		flag.Float64Var(varFloat, shortFlag, defaultValue, "")
	}
	*options = append(*options, FlagStruct{
		LongFlag:    longFlag,
		ShortFlag:   shortFlag,
		Description: description})
}

func appendBoolean(options *[]FlagStruct, varBoolean *bool, longFlag string, shortFlag string, defaultValue bool, description string) {
	flag.BoolVar(varBoolean, longFlag, defaultValue, "")
	if shortFlag != longFlag {
		flag.BoolVar(varBoolean, shortFlag, defaultValue, "")
	}
	*options = append(*options, FlagStruct{
		LongFlag:    longFlag,
		ShortFlag:   shortFlag,
		Description: description})
}
